
import numpy as np

class LocFeature2ImgIndex():
    
    def __init__(self) -> None:
        #每个图像的特征数量
        self._featureCountArr = []
        #图像ID列表
        self._indexArr = []

        self._feat_to_img_dict = {}
        pass
    
    def addIndices(self,imgid,featurecount):
        index0 = 0
        index1 = featurecount
        if len(self._featureCountArr) > 0:
           index0 = self._featureCountArr[-1]
           index1 = index0 + featurecount
        self._featureCountArr.append(index1)
        self._indexArr.append(imgid)
        for i in range( index0, index1):
            self._feat_to_img_dict[i] = imgid
        return None
    
    def read_h5data(self,h5data):
        self._featureCountArr = h5data["feature_index"].__array__()
        self._indexArr = h5data["image_index"].__array__()
        for imgid_index in range(len(self._indexArr)):
            imgid = self._indexArr[imgid_index]
            index1 = self._featureCountArr[imgid_index]
            index0 =  self._featureCountArr[imgid_index-1] if imgid_index>0 else 0
            for i in range( index0, index1):
                self._feat_to_img_dict[i] = imgid


    #效率太低
    # def search(self,featureid):
    #     for i in range(len(self._featureCountArr)):
    #         if featureid < self._featureCountArr[i]:
    #             return self._indexArr[i]
    #     return -1
    
    #高效实现
    def search(self,featureid):
        #没有做检查
        return self._feat_to_img_dict[featureid]

    def searchArr(self,featureids):
        res=[]
        
        for i in range(len(featureids)):
            id = self.search(featureids[i])
            res.append(id)
        return np.array(res)

    def searchArr2d(self,predictions):
        
        predictions_ids = np.array(predictions,copy=True)
        size = predictions.shape
        for i in range(size[0]):
            for j in range(size[1]):
                id = predictions[i,j]
                res = self.search(id)
                predictions_ids[i,j] = res
        return predictions_ids.astype(np.uint16)

    def most_common(self,predictions,topk):
        keys = np.unique(predictions)
        dict_k_v = dict()
        for k in keys:
            res = np.where(predictions == k)
            dict_k_v[k] = len(res[0])

        def get_sorted_list(d, reverse=False):
            return sorted(d.items(), key=lambda x:x[1], reverse=reverse)

        d_list = get_sorted_list(dict_k_v,True)
        vv = np.array([item[0] for item in d_list[0:topk]])
        return vv